from traceback import print_stack
from venv import create

import backtrader as bt
import pandas as pd
from backtrader.feeds import PandasData
from datetime import datetime, timedelta
import warnings
import quantstats as qs
import argparse
import datetime
import os

# 设置matplotlib后端为不需要GUI的backend
import matplotlib

matplotlib.use('Agg')  # 使用Agg后端，不需要显示界面
import matplotlib.pyplot as plt
from feeds.akshare_feed import stock_symbols, etf_symbols


class DualThrustStrategy(bt.Strategy):
    params = (
        ('n', 1),
        ('k1', 0.1),
        ('k2', 0.0),
        ('printlog', True),  # comma is required
    )

    def __init__(self):
        self.order = None
        self.buyprice = None
        self.buycomm = None
        self.bar_executed = None
        self.val_start = None

        self.sma3 = bt.indicators.SMA(self.data, period=3)
        self.sma5 = bt.indicators.SMA(self.data, period=5)
        self.sma10 = bt.indicators.SMA(self.data, period=10)
        # self.sma20 = bt.indicators.SMA(self.data, period=20)

    def log(self, txt, dt=None, doprint=False):
        ''' Logging function fot this strategy'''
        if self.p.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print('%s %s' % (dt.isoformat(), txt))

    def start(self):
        self.val_start = self.broker.get_cash()  # keep the starting cash

    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        self.log('OPERATION PROFIT, GROSS %.2f, NET %.2f' %
                 (trade.pnl, trade.pnlcomm))

    def notify_order(self, order):
        curr_bar = f"open: {self.data.open[0]:.2f}, close: {self.data.close[0]:.2f}, high: {self.data.high[0]:.2f}, low: {self.data.low[0]:.2f}"
        if order.status in [order.Submitted, order.Accepted]:
            # Buy/Sell order submitted/accepted to/by broker - Nothing to do
            # if order.status  == order.Accepted:
            #     self.log(
            #         f'{"buy" if order.isbuy() else "sell"} order Accepted, Price: %.2f, Size: %.0f, curr_bar: {curr_bar}' %
            #         (order.created.price,
            #             order.created.size))
            return

        # self.log(curr_bar)
        # Check if an order has been completed
        # Attention: broker could reject order if not enough cash
        if order.status in [order.Completed]:  # order.Partial
            if order.isbuy():
                self.log(
                    f'BUY EXECUTED, Price: %.2f, Size: %.0f, Cost: %.2f, Comm %.2f, RemSize: %.0f, RemCash: %.2f, curr_bar: {curr_bar}'
                    % (order.executed.price, order.executed.size,
                       order.executed.value, order.executed.comm,
                       order.executed.remsize, self.broker.get_cash()))

                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            else:  # Sell
                self.log(
                    f'SELL EXECUTED, Price: %.2f, Size: %.0f, Cost: %.2f, Comm %.2f, RemSize: %.0f, RemCash: %.2f, curr_bar: {curr_bar}'
                    % (order.executed.price, order.executed.size,
                       order.executed.value, order.executed.comm,
                       order.executed.remsize, self.broker.get_cash()))

            self.bar_executed = len(self)
        elif order.status in [
                order.Canceled, order.Expired, order.Margin, order.Rejected
        ]:
            self.log(
                f'{"buy" if order.isbuy() else "sell"} Order Failed, curr_bar: {curr_bar}'
            )

        self.order = None

    def next(self):
        # Simply log the closing price of the series from the reference
        # self.log('Close, %.2f' % self.data.close[0])
        if self.order:
            return

        # need n day trading range
        if len(self.datas[0]) <= self.p.n:
            return

        highs = self.data.high.get(0, self.p.n)
        lows = self.data.low.get(0, self.p.n)
        opens = self.data.open.get(0, self.p.n)
        closes = self.data.close.get(0, self.p.n)

        h_h, h_c, l_c, l_l = max(highs), max(closes), min(closes), min(lows)
        signal_range = max(h_h - l_c, h_c - l_l)

        current_open = self.data.open[0]
        buytrig = current_open + self.p.k1 * signal_range
        selltrig = current_open - self.p.k2 * signal_range

        sma_min = min(self.sma3[0], self.sma5[0], self.sma10[0])
        sma_max = max(self.sma3[0], self.sma5[0], self.sma10[0])

        prev_sma_diff = self.sma3[-1] - self.sma10[-1]
        sma_diff = self.sma3[0] - self.sma10[0]

        # current_price = self.data.close[0]
        current_price = self.data.close[0]
        self.log(
            f"Buy Trigger: {buytrig:.2f}, Sell Trigger: {selltrig:.2f}, Current Price: {current_price:.2f}"
        )

        if current_price > buytrig:  # buy on upper break
            if self.position.size > 0:
                return
            # if current_price <= self.sma5[0]:
            #     return
            # if self.sma3[0] <= self.sma5[0]:
            #     return

            if prev_sma_diff - sma_diff >= 0 or sma_diff <= 0:
                return

            target = int(self.broker.get_cash() / current_price / 10) * 10
            self.order = self.order_target_size(target=target)
            # self.order = self.order_target_size(target=target, exectype=bt.Order.Limit, price=current_price)
            self.log(
                'BUY ORDER SENT, Pre-pre-close: %.2f, pre-Price: %.2f, buy-price: %0.2f, UB: %.2f, Target size: %.2f'
                % (closes[-1], closes[0], self.order.created.price, buytrig,
                   target))
        elif current_price < selltrig:  # sell on down break
            if self.position.size <= 0:
                return

            # if current_price > self.sma3[0]:
            #     return

            # if min(closes[-1], opens[-1]) < min(closes[0], opens[0]) and \
            #   max(closes[-1], opens[-1]) > max(closes[0], opens[0]):
            #     return

            if self.sma3[0] >= self.sma5[0]:
                return

            # target = -int(self.broker.get_value() / current_price * 0.95)
            target = 0
            self.order = self.order_target_size(target=target)
            # self.order = self.order_target_size(target=target, exectype=bt.Order.Limit, price=current_price)
            self.log(
                'SELL ORDER SENT, Pre-Price: %.2f, Price: %.2f, sell-price: %.2f, LB: %.2f, Target size: %.2f'
                % (closes[-1], closes[0], self.order.created.price, selltrig,
                   target))

    def stop(self):
        # calculate the actual returns
        print(self.analyzers)
        roi = (self.broker.get_value() / self.val_start) - 1.0
        self.log('ROI:        {:.2f}%'.format(100.0 * roi))
        self.log('(Dual thrust params (%2d, %.2f, %.2f)) Ending Value %.2f' %
                 (self.p.n, self.p.k1, self.p.k2, self.broker.getvalue()),
                 doprint=True)


def runstrat(datapath=None, args=None):
    args = parse_args(args)

    cerebro = bt.Cerebro()  # 创建Cerebro实例

    # Get a pandas dataframe
    # datapath = '../../datas/603893.SH_20200207_20250915.csv'
    # datapath = '../../datas/688981.SH_20200716_20250915.csv'
    # datapath = '../../datas/603290.SH_20200204_20250915.csv'
    # datapath = '../../datas/002049.SZ_20050606_20250915.csv'
    # datapath = '../../datas/yhoo-1996-2014.txt'

    # Simulate the header row isn't there if noheaders requested
    dkwargs = dict()
    if args.fromdate is not None:
        fromdate = datetime.datetime.strptime(args.fromdate, '%Y-%m-%d')
        fromdate = datetime.datetime.strptime('2025-01-01', '%Y-%m-%d')
        # fromdate = datetime.datetime.strptime('2020-01-01', '%Y-%m-%d')
        dkwargs['fromdate'] = fromdate

    if args.todate is not None:
        todate = datetime.datetime.strptime(args.todate, '%Y-%m-%d')
        dkwargs['todate'] = todate

    df = pd.read_csv(datapath)
    df.index = pd.to_datetime(df["date"])
    df['openinterest'] = 1000000
    df = df[['open', 'high', "low", "close", "volume", 'openinterest']]
    print(df.head())
    data0 = bt.feeds.PandasDirectData(dataname=df)
    # if dataset is None, args.data has been given
    # data0 = bt.feeds.BacktraderCSVData(dataname=datapath, **dkwargs)
    cerebro.adddata(data0)

    startcash = 20000
    cerebro.addstrategy(DualThrustStrategy)

    # 设置资金
    cerebro.broker.setcash(startcash)
    # 佣金
    cerebro.broker.setcommission(commission=0.0001)
    cerebro.addanalyzer(bt.analyzers.PyFolio, _name="PyFolio")
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='my_sharpe')
    cerebro.addanalyzer(bt.analyzers.Returns, _name='my_returns')
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='my_drawdown')
    cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name='my_trade_analyzer')

    results = cerebro.run()

    sharpe_ratio = results[0].analyzers.my_sharpe.get_analysis()['sharperatio']
    annual_return = results[0].analyzers.my_returns.get_analysis()['rnorm']
    max_drawdown = results[0].analyzers.my_drawdown.get_analysis(
    )["max"]["drawdown"] / 100
    trade_num = results[0].analyzers.my_trade_analyzer.get_analysis(
    )['total']['total']
    #     value_df = pd.DataFrame([results[0].analyzers.my_value.get_analysis()]).T
    #     value_df.columns=['value']
    #     value_df['datetime']=pd.to_datetime(value_df.index)
    #     value_df['date']=[i.date() for i in value_df['datetime']]
    #     value_df= value_df.drop_duplicates("date",keep="last")
    #     value_df = value_df[['value']]
    #     value_df.to_csv("ETF策略的测试结果.csv")
    print(f"夏普率：{sharpe_ratio:.3f}, 年化收益率:{annual_return:.3f}, " +
          f"最大回撤：{max_drawdown:.3f}, 交易次数:{trade_num}")

    #获取回测结束后的总资金
    portvalue = cerebro.broker.getvalue()
    pnl = portvalue - startcash

    #打印结果
    print(f'回测数据: {datapath}')
    print(f'初始资金：{startcash}')
    print(f'总资金: {round(portvalue,2)}')
    print(f'净收益: {round(pnl,2)}')

    # plot result - 在无头环境中安全地绘图
    # try:
    #     cerebro.plot(
    #         style='candle',
    #         barup='red',  # 上涨K线 - 红色
    #         bardown='green',  # 下跌K线 - 绿色
    #         volup='#ffcccc',  # 上涨成交量 - 淡红色
    #         voldown='#ccffcc',  # 下跌成交量 - 淡绿色
    #         # figsize=(12, 8),
    #         # plotdist=0.1
    #         iplot=False,  # 设置为False以避免交互式显示
    #     )
    # except Exception as e:
    #     print(f"绘图失败，跳过绘图步骤: {e}")
    #     print("这通常在无头环境中是正常的")

    # Quantstats分析 - 使用安全的输出方式避免窗口闪退
    pyfolio = results[0].analyzers.PyFolio
    returns, positions, transactions, gross_lev = pyfolio.get_pf_items()

    # 将报告输出到文件而不是弹出窗口
    try:
        # 生成HTML报告文件
        qs.reports.html(returns,
                        output='quantstats_report.html',
                        title='Dual Thrust Strategy Report')
        print("Quantstats报告已保存为 quantstats_report.html")

        # 输出基本统计信息到控制台
        print("\n=== 策略表现统计 ===")
        # 安全地获取和格式化统计数据，处理可能返回Series的情况
        try:
            # 使用.item()方法从Series中提取标量值
            total_return = qs.stats.comp(returns)
            total_return = total_return.item() if hasattr(
                total_return, 'item') else total_return
            print(f"总收益率: {total_return:.2%}")

            annual_return = qs.stats.cagr(returns)
            annual_return = annual_return.item() if hasattr(
                annual_return, 'item') else annual_return
            print(f"年化收益率: {annual_return:.2%}")

            max_dd = qs.stats.max_drawdown(returns)
            max_dd = max_dd.item() if hasattr(max_dd, 'item') else max_dd
            print(f"最大回撤: {max_dd:.2%}")

            sharpe = qs.stats.sharpe(returns, 0.01)
            sharpe = sharpe.item() if hasattr(sharpe, 'item') else sharpe
            print(f"夏普比率: {sharpe:.2f}")

            volatility = qs.stats.volatility(returns)
            volatility = volatility.item() if hasattr(volatility,
                                                      'item') else volatility
            print(f"波动率: {volatility:.2%}")

            win_rate = qs.stats.win_rate(returns)
            win_rate = win_rate.item() if hasattr(win_rate,
                                                  'item') else win_rate
            print(f"胜率: {win_rate:.2%}")
        except Exception as inner_e:
            print(f"统计数据格式化失败: {inner_e}")
            print("使用替代方式显示统计数据:")
            # 替代显示方式，直接打印原始值
            print(f"总收益率: {qs.stats.comp(returns)}")
            print(f"年化收益率: {qs.stats.cagr(returns)}")
            print(f"最大回撤: {qs.stats.max_drawdown(returns)}")
            print(f"夏普比率: {qs.stats.sharpe(returns)}")
            print(f"波动率: {qs.stats.volatility(returns)}")
            print(f"胜率: {qs.stats.win_rate(returns)}")

    except Exception as e:
        print(f"Quantstats分析失败: {e}")
        print_stack()
        print("跳过Quantstats分析")


def parse_args(pargs=None):
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description='Sample for DualThrust example')

    parser.add_argument('--fromdate',
                        required=False,
                        default='2005-01-01',
                        help='Starting date in YYYY-MM-DD format')

    parser.add_argument('--todate',
                        required=False,
                        default=None,
                        help='Ending date in YYYY-MM-DD format')

    parser.add_argument('--cash',
                        required=False,
                        action='store',
                        type=float,
                        default=50000,
                        help=('Cash to start with'))

    parser.add_argument('--cashalloc',
                        required=False,
                        action='store',
                        type=float,
                        default=0.20,
                        help=('Perc (abs) of cash to allocate for ops'))

    parser.add_argument('--commperc',
                        required=False,
                        action='store',
                        type=float,
                        default=0.0033,
                        help=('Perc (abs) commision in each operation. '
                              '0.001 -> 0.1%%, 0.01 -> 1%%'))

    parser.add_argument('--macd1',
                        required=False,
                        action='store',
                        type=int,
                        default=12,
                        help=('MACD Period 1 value'))

    parser.add_argument('--macd2',
                        required=False,
                        action='store',
                        type=int,
                        default=26,
                        help=('MACD Period 2 value'))

    parser.add_argument('--macdsig',
                        required=False,
                        action='store',
                        type=int,
                        default=9,
                        help=('MACD Signal Period value'))

    parser.add_argument('--atrperiod',
                        required=False,
                        action='store',
                        type=int,
                        default=14,
                        help=('ATR Period To Consider'))

    parser.add_argument('--atrdist',
                        required=False,
                        action='store',
                        type=float,
                        default=3.0,
                        help=('ATR Factor for stop price calculation'))

    parser.add_argument('--smaperiod',
                        required=False,
                        action='store',
                        type=int,
                        default=30,
                        help=('Period for the moving average'))

    parser.add_argument('--dirperiod',
                        required=False,
                        action='store',
                        type=int,
                        default=10,
                        help=('Period for SMA direction calculation'))

    parser.add_argument('--riskfreerate',
                        required=False,
                        action='store',
                        type=float,
                        default=0.01,
                        help=('Risk free rate in Perc (abs) of the asset for '
                              'the Sharpe Ratio'))
    # Plot options
    parser.add_argument('--plot',
                        '-p',
                        nargs='?',
                        required=False,
                        metavar='kwargs',
                        const=True,
                        help=('Plot the read data applying any kwargs passed\n'
                              '\n'
                              'For example:\n'
                              '\n'
                              '  --plot style="candle" (to plot candles)\n'))

    if pargs is not None:
        return parser.parse_args(pargs)

    return parser.parse_args()


def selectBtStock():
    stock_data_dir = os.path.abspath(
        os.path.join(os.path.dirname(__file__), '../../', 'datas'))
    symbols = stock_symbols
    symbols.update(etf_symbols)
    index = 0
    print("请选择要回测的股票：")
    for stock_code, sotck_name in symbols.items():
        print(f'{index}: {stock_code}_{sotck_name}')
        index += 1

    selectI = 0
    try:
        # 等待输入index
        selectI = int(input("请输入index: "))
        if selectI < 0 or selectI >= index:
            print("输入错误, 直接退出")
            return
        stock_key = list(symbols.keys())[selectI]
        print(f'选择股票: {stock_key}')
        # 在../../datas/目录查找f'{stock_key}_'开头的行情文件

        return [
            os.path.join(stock_data_dir, f) for f in os.listdir(stock_data_dir)
            if f.startswith(f'{stock_key}_')
        ][0]
    except:
        print("输入错误, 直接退出")
        return


if __name__ == '__main__':
    stock_data = selectBtStock()
    print("回测股票历史行情文件: ", stock_data)
    runstrat(stock_data)
